#!/usr/bin/python
# -*- coding: utf-8 -*-

# External imports
import itertools
import random

# Internal imports
import cryptogram


morsetab = {
  'A': '.-',   'a': '.-',
  'B': '-...', 'b': '-...',
  'C': '-.-.', 'c': '-.-.',
  'D': '-..',  'd': '-..',
  'E': '.',    'e': '.',
  'F': '..-.', 'f': '..-.',
  'G': '--.',  'g': '--.',
  'H': '....', 'h': '....',
  'I': '..',   'i': '..',
  'J': '.---', 'j': '.---',
  'K': '-.-',  'k': '-.-',
  'L': '.-..', 'l': '.-..',
  'M': '--',   'm': '--',
  'N': '-.',   'n': '-.',
  'O': '---',  'o': '---',
  'P': '.--.', 'p': '.--.',
  'Q': '--.-', 'q': '--.-',
  'R': '.-.',  'r': '.-.',
  'S': '...',  's': '...',
  'T': '-',    't': '-',
  'U': '..-',  'u': '..-',
  'V': '...-', 'v': '...-',
  'W': '.--',  'w': '.--',
  'X': '-..-', 'x': '-..-',
  'Y': '-.--', 'y': '-.--',
  'Z': '--..', 'z': '--..',
  '0': '-----',
  '1': '.----',
  '2': '..---',
  '3': '...--',
  '4': '....-',
  '5': '.....',
  '6': '-....',
  '7': '--...',
  '8': '---..',
  '9': '----.',
  ',': '--..--',
  '.': '.-.-.-',
  '?': '..--..',
  ';': '-.-.-.',
  ':': '---...',
  "'": '.----.',
  '-': '-....-',
  '/': '-..-.',
  '(': '-.--.-',
  ')': '-.--.-',
  '_': '..--.-',
  '=': '-...-',
  ' ': '   ',
  '_': ''
}


class MethodError(Exception):
  pass


class Method(object):
  pass


class AntiVigenere(Method):

  @staticmethod
  def decrypt(cryptogram, alphabet):
    if not (cryptogram.hasMsg() and cryptogram.hasKey() and
        cryptogram.hasPlain()):
      raise MethodError, 'Incomplete cryptogram'
    
    # Let the decryption ends!!!
    msg = map(alphabet.encode, cryptogram.msg)
    cryptogram.msg = ''
    plain = map(alphabet.encode, cryptogram.plain)
    cryptogram.plain = ''
    for i in xrange(len(msg)):
      cryptogram.msg += alphabet.decode(msg[i] - msg[(i +
        len(cryptogram.key)) % len(msg)])
      cryptogram.plain += alphabet.decode(plain[i] - plain[(i +
        len(cryptogram.key)) % len(plain)])


class AntiJPS(Method):

  @staticmethod
  def decrypt(cryptogram):
    if not (cryptogram.hasMsg() and
        cryptogram.hasPlain()):
      raise MethodError, 'Incomplete cryptogram'


class Morse(Method):

  @staticmethod
  def crypt(cryptogram):
    if not cryptogram.hasPlain():
      raise MethodError, 'Incomplete cryptogram'
    msg = ''
    for c in cryptogram.plain:
      msg += '%s ' % morsetab.get(c, '@')
    cryptogram.msg = msg


  @staticmethod
  def decrypt(cryptogram):
    if not cryptogram.hasMsg():
      raise MethodError, 'Incomplete cryptogram'
    plain = ''
    for t in cryptogram.msg.split(' '):
      for c, m in morsetab.iteritems():
        if m == t:
          break
      else:
        # print 'unknown', t
        c = '@'
      plain += c
    cryptogram.plain = plain


class Llamame36(Method):

  @staticmethod
  def getAlphabets(key):
    CIH = cryptogram.Alphabet(key)
    CIH.compress(cryptogram.Llamame36Alphabet())
    CIV = cryptogram.Alphabet(key)
    CIV.compress(cryptogram.Llamame36Alphabet(), preExist=False,
        lastOrder=-1)
    CIM = cryptogram.Alphabet(key)
    CIM.compress(cryptogram.Llamame36Alphabet(), preExist=False,
        existOrder=-1)
    return CIH, CIV, CIM

  @staticmethod
  def crypt(cryptogram, CIH, CIV, CIM):
    if not (cryptogram.hasPlain() and cryptogram.hasKey()):
      raise MethodError, 'Incomplete cryptogram'

    # Let the encryption begin!!!
    msg = []
    for c in cryptogram.plain:
      index = CIM.encode(c)
      row = CIV.decode(index / 6 * 6 + random.randrange(6))
      column = CIH.decode((index % 6) * 6 + random.randrange(6))
      msg.append('%c%c' % (row, column))
    cryptogram.msg = ' '.join(msg)

  @staticmethod
  def decrypt(cryptogram, CIH, CIV, CIM):
    if not (cryptogram.hasMsg() and cryptogram.hasKey()):
      raise MethodError, 'Incomplete cryptogram'
    
    # Let the decryption ends!!!
    msg = ''
    for c in cryptogram.msg.split():
      row, column = c[0], c[1]
      msg += CIM.decode(CIV.encode(row) / 6 * 6 +
          CIH.encode(column) / 6)
    cryptogram.plain = msg


# TODO: Llamame36B only changes the fact that CIM rotates with spaces.

class Vigenere(Method):

  @staticmethod
  def crypt(cryptogram, alphabet):
    if not (cryptogram.hasPlain() and cryptogram.hasKey()):
      raise MethodError, 'Incomplete cryptogram'
    
    # Let the encryption begin!!!
    msg = ''
    for c, d in itertools.izip(cryptogram.plain,
        itertools.cycle(cryptogram.key)):
      msg += alphabet.decode(alphabet.encode(c) + alphabet.encode(d))
    cryptogram.msg = msg

  @staticmethod
  def decrypt(cryptogram, alphabet):
    if not (cryptogram.hasMsg() and cryptogram.hasKey()):
      raise MethodError, 'Incomplete cryptogram'
    
    # Let the decryption ends!!!
    msg = ''
    for c, d in itertools.izip(cryptogram.msg,
        itertools.cycle(cryptogram.key)):
      msg += alphabet.decode(alphabet.encode(c) - alphabet.encode(d))
    cryptogram.plain = msg

class EnigmaII(Method):

  @staticmethod
  def crypt(cryptogram, alphabet):
    if not (cryptogram.hasPlain() and cryptogram.hasKey()):
      raise MethodError, 'Incomplete cryptogram'

    # Let the encryption begin!!!
    msg = ''
    keyLength = len(cryptogram.key)
    J = (keyLength % 2) and -1 or 1
    idx = 0
    cdx = 1
    for c, d in itertools.izip(cryptogram.plain,
        itertools.cycle(cryptogram.key)):
      if cdx % 2: # Impar
        A = (idx + 1)
        B = alphabet.encode(c, True) + 1
        C = alphabet.encode(d, A % 2 != 0) + 1
        D = B + J * A - J * C
        # print A, B, C, D, alphabet.decode(D - 1, True)
        msg += alphabet.decode(D - 1, True)
      else: # Par
        A = (idx + 1)
        B = alphabet.encode(c, False) + 1
        C = alphabet.encode(d, A % 2 != 0) + 1
        D = B - J * A + J * C
        # print A, B, C, D, alphabet.decode(D - 1, False)
        msg += alphabet.decode(D - 1, False)
      idx = (idx + 1) % keyLength
      cdx += 1
    cryptogram.msg = msg

  @staticmethod
  def decrypt(cryptogram, alphabet):
    if not (cryptogram.hasMsg() and cryptogram.hasKey()):
      raise MethodError, 'Incomplete cryptogram'

    # Let the encryption begin!!!
    msg = ''
    keyLength = len(cryptogram.key)
    J = (keyLength % 2) and -1 or 1
    idx = 0
    cdx = 1
    for c, d in itertools.izip(cryptogram.msg,
        itertools.cycle(cryptogram.key)):
      if cdx % 2: # Impar
        A = (idx + 1)
        B = alphabet.encode(c, True) + 1
        C = alphabet.encode(d, A % 2 != 0) + 1
        D = B - J * A + J * C
        # print A, B, C, D, alphabet.decode(D - 1, True)
        msg += alphabet.decode(D - 1, True)
      else: # Par
        A = (idx + 1)
        B = alphabet.encode(c, False) + 1
        C = alphabet.encode(d, A % 2 != 0) + 1
        D = B + J * A - J * C
        # print A, B, C, D, alphabet.decode(D - 1, False)
        msg += alphabet.decode(D - 1, False)
      idx = (idx + 1) % keyLength
      cdx += 1
    cryptogram.plain = msg


class DASK(Method):

  @staticmethod
  def dec2bin(n):
    s = ''
    while n:
      s = ((n % 2) and '1' or '0') + s
      n = n >> 1
    while len(s) < 5:
      s = '0' + s
    return s

  @staticmethod
  def crypt(cryptogram, alphabet):
    if not (cryptogram.hasPlain() and cryptogram.hasKey()):
      raise MethodError, 'Incomplete cryptogram'

    # Let the encryption begin!!!
    while len(cryptogram.plain) % len(cryptogram.key) != 0:
      cryptogram.plain += alphabet.randChoice()
    binaryPlain = ''.join([DASK.dec2bin(alphabet.encode(c))
      for c in cryptogram.plain])
    binaryKey = ''.join([DASK.dec2bin(alphabet.encode(c))
      for c in cryptogram.key])
    keyVowels = len([c
      for c in cryptogram.key
      if c in ('A', 'E', 'I', 'O', 'U')])
    rotatedPlain = ''
    while binaryPlain:
      block = binaryPlain[:len(cryptogram.key) * 5]
      binaryPlain = binaryPlain[len(cryptogram.key) * 5:]
      rotatedPlain += block[-keyVowels:] + block[:-keyVowels]
    binaryMsg = ''
    for c, d in itertools.izip(rotatedPlain, itertools.cycle(binaryKey)):
      binaryMsg += str(int(c, 2) ^ int(d, 2))
    msg = ''
    while binaryMsg:
      block, binaryMsg = binaryMsg[:5], binaryMsg[5:]
      msg += alphabet.decode(int(block, 2))
    cryptogram.msg = msg

  @staticmethod
  def decrypt(cryptogram, alphabet, vowels=0):
    if not (cryptogram.hasMsg() and cryptogram.hasKey()):
      raise MethodError, 'Incomplete cryptogram'

    # Let the encryption begin!!!
    binaryMsg = ''.join([DASK.dec2bin(alphabet.encode(c))
      for c in cryptogram.msg])
    binaryKey = ''.join([DASK.dec2bin(alphabet.encode(c))
      for c in cryptogram.key])
    keyVowels = vowels or len([c
      for c in cryptogram.key
      if c in ('A', 'E', 'I', 'O', 'U')])
    rotatedPlain = ''
    for c, d in itertools.izip(binaryMsg, itertools.cycle(binaryKey)):
      rotatedPlain += str(int(c, 2) ^ int(d, 2))
    binaryPlain = ''
    while rotatedPlain:
      block = rotatedPlain[:len(cryptogram.key) * 5]
      rotatedPlain = rotatedPlain[len(cryptogram.key) * 5:]
      binaryPlain += block[keyVowels:] + block[:keyVowels]
    plain = ''
    while binaryPlain:
      block, binaryPlain = binaryPlain[:5], binaryPlain[5:]
      plain += alphabet.decode(int(block, 2))
    cryptogram.plain = plain

